Scaling MLPs: A Tale of Inductive Bias — part 3/4
Cleaned-text code points [28082, 43904). Bibliography omitted here; source.md includes it.
Following OpenAI, (2018) we define the computational cost C C incurred from training a model f f on N N examples for T T epochs as C = FLOP ( f ) × 3 × N × T , C=\text{FLOP}(f)\times 3\times N\times T, (4) where FLOP ( f ) \text{FLOP}(f) denotes the number of FLOPs needed to complete the forward pass of f f for a single example. We note that the number of parameters P P present in f f enters this equation implicitly in the form of FLOP ( f ) ∝ P \text{FLOP}(f)\propto P . Observe that a given level of compute can be achieved in different ways, i.e. using more parameters P P , training on more examples N N , or training for a longer time T T . When allocating a given level of compute optimally, it is observed that for convolutional and transformer-based architectures, the test error E ( C ) E(C) as a function of compute behaves as a power-law E ( C ) = a ( b + C ) − α + E ∞ , E(C)=a(b+C)^{-\alpha}+E_{\infty}, (5) where a , b , E ∞ ∈ ℝ + a,b,E_{\infty}\in\mathbb{R}_{+} and α > 0 \alpha>0 is the scaling coefficient determining the rate of decay. E ∞ E_{\infty} denotes the irreducible error, i.e. even if infinite compute were employed, the performance remains imperfect. The test error can be measured upstream (i.e. on the pre-training task) or downstream when fine-tuning on a different task. We investigate various pre-training schemes with different number of examples, parameter counts and training times. We subsample ImageNet21k proportionally across classes and pre-train variously sized inverted bottleneck MLPs. We summarize the configurations in Table 5 . We then measure test error on the downstream task of CIFAR100 in Fig. 1 as well as CIFAR10 and ImageNet1k in Fig. 6 by linearly transferring the learnt features (without test-time augmentations). The plotting style is inspired by Zhai et al., (2022) . Each point in the curve is the downstream performance of an MLP, where the color of the point indicates the model type (blue denotes smaller and red larger models) and the size of the point indicates the number of pre-training examples. Points connected by a line indicates longer training times where T ∈ { 50,100,200,400,800 } T\in\{50,100,200,400,800\} is measured in epochs. In all experiments, we employ data augmentation for pre-training. We observe that the compute-optimal performance of MLPs strongly exhibits characteristics of a power-law with coefficients α ∈ { 0.12 , 0.25 , 0.35 } \alpha\in\{0.12,0.25,0.35\} . This is very encouraging for future theoretical work, showing that MLPs indeed mirror the scaling behaviour of modern models. We provide the dual results for the standard MLPs in Appendix B.2 , noting that they exhibit essentially the same scaling behaviour, albeit with a slightly weaker slope and intercept. We further study how performance E E evolves when compute is either bottlenecked by the number of parameters P P or the dataset size N N . We visualize the resulting scaling laws in Fig. 7 . We find a very steep decay rate in terms of parameters P P where roughly α P ≈ 1 \alpha_{P}\approx 1 , whereas for dataset size N N we identify a significantly slower rate of α N ≈ 0.35 \alpha_{N}\approx 0.35 . This shows that the performance of MLPs is significantly more limited by the dataset size, which is in-line with the fact that MLPs exhibit a bad inductive bias. We investigate the role of dataset size and parameters more in the next paragraph.
#### Parameters or examples.
Given a fixed level of compute C C , what is the optimal way to allocate it to parameter count P P and number of examples N N ? In order to be more comparable to previous work, we assume a fixed training time T = 50 T=50 . To answer this question, we follow the approach outlined in Hoffmann et al., (2022) and plot the optimal compute models identified in Fig. 1 both against model size P P and number of examples N N . We visualize the results in Fig. 8 . We empirically observe that the optimal parameter count P ∗ ( C ) P^{*}(C) and dataset size N ∗ ( C ) N^{*}(C) as a function of compute C C exhibit power-law behaviour of the approximate form
P ∗ ( C ) ∝ C 0.35 N ∗ ( C ) ∝ C 0.65 P^{*}(C)\propto C^{0.35}\hskip 14.22636ptN^{*}(C)\propto C^{0.65} While for transformers, the number of examples (or tokens) N N and parameters P P are scaled equally ( Hoffmann et al.,, 2022 ) (i.e. α P ≈ α N ≈ 0.5 \alpha_{P}\approx\alpha_{N}\approx 0.5 ), in contrast we observe that the optimal strategy for MLPs invests significantly more compute into dataset size N N . This is further evidence for the weaker inductive bias present in MLPs, which needs more examples in order to be compensated for.
### 4.5 Computational Feasibility
We believe that a further exciting feature of our study is its computational feasibility, while at the same time preserving the main characteristics of large-scale pre-training. All of our experiments were conducted on a single NVIDIA RTX A5000 GPU with 24GB of memory. In conjunction with the strongly optimized FFCV dataloading framework ( Leclerc et al.,, 2023 ) and the inherent efficiency of MLPs, we are able to perform very rapid training. For instance we complete a single epoch on ImageNet21k with the B-12/Wi-1024 architecture, equipped with 124 124 million parameters, in only roughly 450 450 seconds, while the smaller variant B-6/Wi-1024 at a parameter count of 74 74 million requires roughly 250 250 seconds on the specified hardware. Low memory requirements allow us to train with a batch-size of 16384 16384 without having to shard computation among multiple GPUs. We compare the computational efficiency of MLPs with contemporary networks of similar size such as ResNet-152 , ViT-B/4 and ViT-B/8 in Appendix A.5 .
## 5 Related Works
There are some prior works that investigate MLPs on vision tasks. Lin et al., (2015) study the performance of MLPs on small scale datasets such as CIFAR10. They observe similar improvements when using inverted bottleneck layers but do not study larger-scale setups, transfer-learning nor do they discuss the implications for theoretical works. The bottleneck structure used in this work has also been investigated theoretically ( Parhi and Nowak,, 2021 ; Shenouda et al.,, 2023 ; Parkinson et al.,, 2023 ) , further highlighting that such an architecture exhibits desirable properties. Urban et al., (2017) study to what degree convolutions are necessary for good performance and conclude that even with distillation techniques it remains very difficult to train performant MLPs on CIFAR10. Other approaches have focused on sparsifying fully-connected layers through evolutionary training ( Mocanu et al.,, 2018 ; Fernando et al.,, 2016 ) , aiming to learn a good inductive bias from scratch. Similarly, Neyshabur, (2020) study how the inductive bias of MLPs can be improved by systematically sparsifying them with a LASSO-type algorithm, making them more convolution-like. d'Ascoli et al., (2019) on the other hand first train a convolutional network for a certain duration and then subsequently continue training the network as an MLP (by using the correspondence between CNNs and MLPs highlighted in Sec. 2 ). They show that good performance can be reached if the network was trained long enough as a CNN. In contrast to these works, our goal is not to enhance the inherent inductive bias of MLPs but study whether it can be overcome with enough scale. The advent of the MLP-Mixer ( Tolstikhin et al.,, 2021 ) has led to a series of follow-up work, similarly using MLPs as a patch processor and token mixer ( Touvron et al.,, 2021 ; Chen et al.,, 2022 ; Lian et al.,, 2022 ; Guo et al.,, 2021 ; Liu et al.,, 2021 ) . Again, we remark that these architectures all possess significantly more inductive bias. Finally, we would like to remark that MLPs are successfully used in other areas such as novel view synthesis (e.g. NeRF ( Mildenhall et al.,, 2021 ) ).
## 6 Discussion
In this work, we have explored the limits of the multi-layer perceptron as an architecture for vision tasks. Our study reveals that (1) lack of inductive bias can be compensated by scale and (2) MLPs constitute a (largely) accurate proxy for modern architectures, further cementing their role as the main theoretical object of study. The role of data augmentation and the implicit bias of SGD however strongly differ for MLPs in the setting considered in this work and theoretical works should take this into account. Large-scale pre-training of MLPs proves to be very efficient, enabling researchers with less access to computational resources to study this very exciting line of work. While lack of inductive bias does not prevent MLPs from reaching impressive performance, it leads to an interesting shift in compute-optimality towards more training examples. Subjecting MLPs to even larger amounts of compute similar to Zhai et al., (2022) , especially in the form of more training examples, remains as very interesting future work.
## Appendix A Experimental Details
### A.1 Resources
For all experiments we rely on NVIDIA RTX A5000 GPU with 24GB of memory. Every experiment can be performed on a single GPU. We leverage the FFCV dataloader framework since the transfer time of the data to the GPU becomes the bottleneck in terms of training time in case of MLPs. All of our experiments were performed in PyTorch ( Paszke et al.,, 2019 ) .
### A.2 Additional Ablations
#### Ablations.
We provide some more ablations in Fig. 9 . More specifically, for a (approximate) fixed budget of compute, we investigate different architecture and optimization choices, when pretraining on ImageNet1k and performing linear probing on CIFAR100.
#### Normalization.
We investigate the importance of the normalization method (LayerNorm vs BatchNorm) in more detail in Table 6 . We pre-train two B-MLPs on ImageNet21k with layer normalization and batch normalization and compare the fine-tuning performance on various tasks. We find that the techniques perform similarly, which layer normalisation having a slight edge.
#### Label smoothing.
We further ablate the influence of label smoothing on the downstream performance. We pre-train B-MLPs with varying amounts of label smoothing ( α ∈ { 0.0 , 0.1 , 0.3 } \alpha\in\{0.0,0.1,0.3\} ) and evaluate the resulting down-stream fine-tuning performance. We report the results in Table 7 . While label smoothing does provide some boost in performance, the gains are very modest. Label smoothing is thus helpful but not essential for training MLPs.
#### Architecture.
We make the following observations/recommendations to boost the model’s performance, in line with results reported in the literature ( Liu et al.,, 2022 ) ; (1) replacing ReLUs and GELUs boosts results significantly, (2) adding skip connections every two layers helps with optimization, especially for deeper networks. (3) Using an inverted bottleneck increases performance even more. (4) Using a normalization layer in the PRE-LN configuration helps with optimization and (4) layer normalization leads to significantly better results compared to batch normalization, while also being more stable during training.
#### Optimization.
As discussed in the main text, augmentations are crucial, and disabling them can have a detrimental effect. We also found that clipping gradients, using weight decay and dropout have a small positive effect on downstream performance. Finally, replacing LION ( Chen et al.,, 2023 ) with Adam(W), leads to a decrease in performance.
### A.3 Linear Probing
We showcase the transferability of our MLPs by training a linear classifier on top of the frozen features. For training the linear layer, we use the LION optimizer with a learning rate of η = 0.00001 \eta=0.00001 for 50 50 epochs. We display the results in Table 8 .
We observe very strong down-stream performance even in this more limited setting, highlighting how transferable the features learnt by MLPs are.
### A.4 Scaling Laws
#### Implementation Details.
For the scaling law plots, we trained all the models with a batch-size 16384 16384 and the LION optimizer with a learning rate η = 0.00001 \eta=0.00001 and weight decay of strength 0.001 0.001 . We further use label smoothing of strength 0.3 0.3 . We again use augmentations in the form of random flips and crops as well as MixUp with strength 0.8 0.8 . We rely on the curvefit function from the SciPy library ( Virtanen et al.,, 2020 ) to fit powerlaws of the form E ( C ) = a ( b + C ) − α + E ∞ E(C)=a(b+C)^{-\alpha}+E_{\infty} .
### A.5 Computational Efficiency
We highlight the fact that although MLPs require a lot of training data, inference is extremely efficient from a computational perspective. To illustrate this, we embark on the following comparison; we study inference on 64 × 64 64\times 64 resolution images in an MLP vs other popular vision architectures of similar size and complexity in Table 9 . More specifically, we compare against a ResNet-152 , where we replace the stride in the first convolutional layer and remove the first max-pooling operation to compensate for the smaller image size. We also compare against a base ViT and Mixer model, where we extract patches from 4 × 4 4\times 4 regions in the original image.
As it quickly becomes eminent, MLPs require significantly less FLOPs to make predictions on individual images, in essence utilizing their parameters a lot more methodically. As a result, latency and throughput are significantly better compared to other candidate architectures. We measure throughput using the optimal batch size on an NVIDIA RTX A5000. We highlight, that our MLPs, in contrast to the other architectures are memory bound, meaning that their throughput is determined by the prefetching bandwidth of our GPU. Hardware advancement and specialized architectures could significantly mitigate this effect. Neglecting memory transfer time by propagating the same input through our network gives a further 6 6 -fold increase in the potential throughput.
## Appendix B Results for Standard MLPs
### B.1 Transfer Learning
For completeness we also analyze the transfer performance of standard MLPs when pre-trained on ImageNet21k. We compare a S-MLP of depth 6 6 and width 2048 2048 against a B-MLP of depth 6 6 and width 1024 1024 , ensuring that both models roughly have a parameter count of around ≈ 70 \approx 70 million. We display the results in Table 10 . We observe that even the features of a standard MLP (i.e. without residual connections and bottleneck structure) transfer very well on different downstream task. Tthe inverted-bottleneck MLP however still remains superior.
### B.2 Scaling Laws
We also evaluate the scaling law of standard MLPs by training variously sized models on different subsets of ImageNet21k and subsequently linearly probing the features on CIFAR100. The setting is identical to the one described in 4.4 . We observe that also standard MLPs exhibit power-law behaviour. The slope ( 0.22 0.22 vs 0.25 0.25 ) and the intercept ( 0.18 0.18 vs 0.16 0.16 ) are however worse when compared against the inverted-bottleneck MLP.
## Appendix C Weight visualizations
We visualize the first layer weights 𝑾 ( 1 ) ∈ ℝ 3 w h × m \bm{W}^{(1)}\in\mathbb{R}^{3wh\times m} by reshaping them back to ℝ w × h × 3 × m \mathbb{R}^{w\times h\times 3\times m} . We then produce a ℝ w × h × m \mathbb{R}^{w\times h\times m} representation by taking the maximal value along the channel dimension. We display such visualizations of the first 5 × 5 = 25 5\times 5=25 "filters" for different pre-training sizes, also including the weights at random initialization in Fig. 11 . All models were trained with data augmentation. We observe that filters increasingly develop structure as we increase the dataset size and become more and more localized.