SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ildus Sadrtdinov $^{1*}$ , Dmitrii Pozdeev $^{1,2*}$ , Dmitry Vetrov $^{3}$ , Ekaterina Lobacheva $^{4}$
$^{1}$ HSE University  $^{2}$ Technical University of Munich (TUM)
$^{3}$ Constructor University, Bremen  $^{4}$ Independent researcher isadrtdinov@hse.ru, dmitrii.pozdeev@tum.de
dvetrov@constructor.university, lobacheva.tjulja@gmail.com
Transfer learning and ensembling are two popular techniques for improving the performance and robustness of neural networks. Due to the high cost of pre-training, ensembles of models fine-tuned from a single pre-trained checkpoint are often used in practice. Such models end up in the same basin of the loss landscape, which we call the pre-train basin, and thus have limited diversity. In this work, we show that ensembles trained from a single pre-trained checkpoint may be improved by better exploring the pre-train basin, however, leaving the basin results in losing the benefits of transfer learning and in degradation of the ensemble quality. Based on the analysis of existing exploration methods, we propose a more effective modification of the Snapshot Ensembles (SSE) for transfer learning setup, StarSSE, which results in stronger ensembles and uniform model soups.
Ensembling neural networks by averaging their predictions is one of the standard ways to improve model quality. This technique is known to improve not only the task-specific quality metric, e. g., accuracy, but also the model calibration (Lakshminarayanan et al., 2017) and robustness to noise (Gustafsson et al., 2020), domain shift (Ovadia et al., 2019), and adversarial attacks (Strauss et al., 2017). In practice, the importance of the solution robustness often goes hand in hand with the scarcity of data and computational resources for training. A popular way to obtain a high-quality solution for a task with limited data using less computation is transfer learning, which initializes the training process on the target task with the model checkpoint pre-trained on a large source task.
In this paper, we investigate the effectiveness of neural network ensembling methods in transfer learning setup since combining the two techniques is a non-trivial problem. The diversity of networks is crucial for obtaining substantial quality gain with ensembling. Therefore, in classic deep ensembles (DE by Lakshminarayanan et al. (2017)), networks are trained independently from different random initializations. However, in the transfer learning setup, networks are not trained from scratch but are fine-tuned from pre-trained checkpoints. Fine-tuning networks from different independently pretrained checkpoints results in a diverse set of models and a high-quality final ensemble but requires a huge amount of computation due to the large cost of the pre-training process. Fine-tuning networks from the same pre-trained checkpoint is much more resource-effective; however, the resulting set of networks are usually located in the same basin of the loss landscape (Neyshabur et al., 2020), are less diverse, and their ensemble has lower quality (Mustafa et al., 2020). So, the main question of our paper is the following: can we improve the quality of the ensemble of networks obtained by fine-tuning from the same pre-trained checkpoint by better exploring the loss landscape of the target task?
We focus on the image classification target tasks and consider two types of pre-training: supervised (Kornblith et al., 2019) and self-supervised (Grill et al., 2020). First of all, in Section 5.1, we confirm that in this setup, there is indeed a substantial quality gap between ensembles of models fine-tuned from the same and different pre-trained checkpoints. Thus, methods aiming at closing this gap are worth researching. We also show that models fine-tuned from the same pre-trained checkpoint lay in the same basin of the loss landscape, meaning there is no significant high-loss barrier on the linear interpolation between them. We refer to this basin as the pre-train basin throughout the paper.
To understand whether the ensemble quality can be improved by better exploring the pre-train basin or the close vicinity outside of it, we take inspiration from the literature on efficient ensembling methods in regular non-transfer learning setup. We focus our study on cyclical learning rate methods, such as Snapshot Ensembling (SSE by Huang et al. (2017)), because they can either explore one basin of the loss landscape or go outside of the basin depending on the chosen hyperparameters. In Section 5.2, we show that with accurate hyperparameter tuning, these methods can explore the pre-train basin and produce ensembles of similar or even higher quality than the independent fine-tuning of several neural networks. However, they are not able to leave the pre-train basin without losing the advantages of transfer learning, followed by significant degradation in the quality of individual models and the whole ensemble. Thus, we conclude that the gap between ensembles of networks fine-tuned from the same and from different independently pre-trained checkpoints can not be closed by exploring the loss landscape of the target task using existing methods.
Based on the analysis of the behavior of the standard sequential variant of SSE in Section 5.3, we propose its parallel variant StarSSE, which is more suitable for exploring the pre-train basin in the transfer learning setup (see Section 5.4). StarSSE fine-tunes the first model from the pre-trained checkpoint in a usual way and then trains all the following models, independently initializing them with the first fine-tuned one. In contrast to the SSE, StarSSE does not move further from the pre-trained checkpoint with each training cycle; therefore, it does not suffer from the progressive degradation of individual model quality. Even though StarSSE is technically very similar to the independent fine-tuning of all the models from the pre-trained checkpoint, it has an important conceptual difference: it separates 1) model fine-tuning from the pre-train checkpoint to the area with high-quality solutions for the target task and 2) exploration of this area for obtaining diverse models. As a result, StarSSE can collect a more diverse set of networks without sacrificing the quality of individual models.
Finally, we analyze different aspects of SSE and StarSSE that can be important for their practical application. Firstly, in Section 6, we show that the set of models, collected by SSE or StarSSE in the pre-train basin, can not only constitute an effective ensemble but also result in a strong uniform model soup (Wortsman et al., 2022; Rame et al., 2022), i.e., a model obtained by weight-averaging all models in the set. Then, in Section 7, we study the robustness of the resulting ensembles and model soups to the test data distribution shifts. Lastly, in Section 8, we confirm that SSE and StarSSE are effective in more practical large-scale experiments with CLIP model (Radford et al., 2021) on the ImageNet classification target task (Russakovsky et al., 2015). In all additional experiments, StarSSE also outperforms the standard SSE.
Our code is available at https://github.com/isadrtdinov/ens-for-transfer.
2 Ensembling methods in regular training
Deep ensemble (DE) is the most common method of neural network ensembling in a regular nontransfer learning setup. A DE consists of several networks of the same architecture trained separately from different random initializations using the same training methods and hyperparameters. Training from different initializations makes neural networks diverse and results in a high-quality ensemble. However, separate training of all the networks requires a lot of resources. From the loss landscape perspective, DE can be called a global method: due to diverse initializations, networks in such an ensemble are independent and come from different basins of low loss (Fort et al., 2019).
To reduce training costs, many methods propose to train the networks jointly or consequentially. Most of these methods can be assigned to one of two groups based on the degree of coverage of different loss landscape basins: local and semi-local. Local methods (SWAG by Maddox et al. (2019), SPRO by Benton et al. (2021), FGE by Garipov et al. (2018), etc.) train the first network and then efficiently sample other ensemble members from the basin of low loss around it. Semi-local methods (SSE
by Huang et al. (2017), cSGLD by Zhang et al. (2020), etc.), instead of focusing on one basin, try to explore a larger portion of the loss landscape by sequentially training networks using a cyclical learning rate schedule. At the beginning of each cycle, the learning rate is increased to help the network leave the current low-loss basin, and at the end of the cycle, it is decreased to obtain a high-quality checkpoint in the next basin to add to the ensemble. Because of the knowledge transfer between the consecutive models, such methods can use training cycles shorter than usual model training, which decreases the overall ensemble training cost. In terms of network diversity and the final ensemble quality, semi-local methods work better than local ones, given the same number of networks in an ensemble. However, methods from both of these groups are inferior to the global ones (Ashukha et al., 2020; Fort et al., 2019).
The described categorization of ensembling methods into three groups, local, semi-local, and global, is not strict. For example, FGE, which is usually assigned to the local group, and SSE, which is usually assigned to the semi-local group, are both based on cyclical learning rate schedules and can be used as either local or semi-local methods, depending on the hyperparameters. Also, local and semi-local methods can be combined with global ones by training several networks from different random initializations, then sampling some networks locally around each of them, and finally ensembling all the obtained networks (Fort et al., 2019). Moreover, recent works on loss landscape analysis hypothesize that two networks trained independently from different random initializations are effectively located in the same basin up to the neuron permutation (Entezari et al., 2022; Ainsworth et al., 2023) and activations renormalization (Jordan et al., 2023), which blurs the boundaries between local and global methods. However, existing local and semi-local methods are not able to find as diverse networks as the global ones can, so in practice, these groups are still easily distinguishable.
If several networks are located in the same basin of low loss, they can be weight-averaged instead of ensembling. The Stochastic Weight Averaging technique (SWA by Izmailov et al. (2018)) averages several checkpoints from one standard training run or models from a local cyclical learning rate schedule method to improve the model quality and robustness.
3 Scope of the paper: transfer learning perspective
In the transfer learning setup, a DE for the target task can be trained in two ways: using different separately pre-trained checkpoints to initialize each network or using the same pre-trained checkpoint for all the networks. In terms of the categorization from the previous section, only the former variant can be called global because the networks fine-tuned from the same pre-trained checkpoints are usually located in the same basin of low loss (Neyshabur et al., 2020). Thus, we call the former variant Global DE while the latter is Local DE. Similarly to the relation between local and global methods in regular training, in the transfer learning setup, Global DEs contain more diverse networks and show better quality than Local DEs (Mustafa et al., 2020). At the same time, the difference in training cost between Local and Global DEs is even more pronounced because the pre-training cost is usually huge compared to the fine-tuning cost.
One of the goals of our paper is to understand whether the existing local and semi-local ensembling methods can be beneficial in the transfer learning setup and help close the quality gap between Local and Global DEs. We focus on the situation when only one pre-trained checkpoint is available, as in Local DE, and study the effectiveness of cyclical learning rate methods. This group of methods contains such methods as FGE, SSE, and cSGLD. We choose cyclical methods because they can work in both local and semi-local regimes and show the strongest results in closing the gap with global methods in the regular training setup (Ashukha et al., 2020). In Appendix C, we show that in the transfer learning setup, they are also more effective than other local methods. Cyclical methods have very similar structure and differ mostly in training scheme for the first network and learning rate schedules in the cycles. FGE trains the first network with a more standard long schedule, while SSE and cSGLD use the same schedule in all the cycles. Also, FGE uses short cycles with a triangular learning rate schedule and low maximum learning rate, while SSE and cSGLD use long cycles with a cosine learning rate schedule and relatively high maximum learning rate. In the paper, we mostly experiment with a version of SSE, which fine-tunes the first network with a fixed standard fine-tuning schedule and then trains all the following networks with the same cosine schedule, for which we vary the length and the maximum learning rate. Depending on the cycle hyperparameters, SSE can behave as either a local method (short cycles with low maximum learning rate) or a semi-local method (longer cycles with higher maximum learning rate). We choose to fine-tune the first network in a
fixed way for easier comparison with Local and Global DEs. The results for other variants of cyclical methods are similar and discussed in Appendix D.
To the best of our knowledge, cyclical methods were not previously studied in the transfer learning setup. Moreover, based on our analysis of the SSE behavior, we propose its parallel version, StarSEE, which shows better results in this setup. Several recent works (Mustafa et al., 2020; Wortsman et al., 2022; Rame et al., 2022) study another technique for improving the diversity of Local DEs by training networks using different hyperparameters (learning rate, weight decay, data augmentation, etc.). In Appendix K, we show that this technique can be also effective for Global DEs and StarSSEs, however, it is not able to close the gaps between the methods and change the comparison results.
Since models in a Local DE are usually located in the same basin of low loss, they can be weight-averaged to obtain a single model called a model soup (Wortsman et al., 2022; Rame et al., 2022). If the set of models to average is chosen properly using validation data, a model soup provides a quality boost without the high computational costs of ensembles in the inference stage. Model soups usually show slightly inferior results on the in-distribution data compared to Local DE, but they are more robust to distributional shifts. We also study the weight-averaging potential of SSE and StarSSE in the transfer learning setup, including the robustness of the resulting models to distributional shifts.
Data, architectures, and pre-training. In most of the experiments, we use a standard ResNet-50 architecture (He et al., 2016) and consider two types of pre-training on the ImageNet dataset (Rusakovsky et al., 2015): supervised and self-supervised with the BYOL method (Grill et al., 2020). We use 5 independently pre-trained checkpoints in both cases. As for target tasks, we choose three image classification tasks: CIFAR-10 (Krizhevsky et al., a), CIFAR-100 (Krizhevsky et al., b), and SUN-397 (Xiao et al., 2010). For SUN-397, we use only the first train/test split of the ten specified in the original paper. We also experiment with a CLIP model (Radford et al., 2021) on the ImageNet target task in Section 8, other non-natural image classification target tasks in Appendix N, and Swin-T transformers (Liu et al., 2021) in Appendix M.
Fine-tuning of individual models. We replace the last fully-connected layer with a randomly initialized layer having an appropriate number of classes and then fine-tune the whole network using mini-batch SGD with batch size 256, momentum 0.9, and cosine learning rate schedule (Loshchilov and Hutter, 2017). For each target task and each type of pre-training, we choose the optimal number of epochs, weight decay, and initial learning rate for fine-tuning by grid search, allocating a validation subset from the training set (see Appendix B for more detail).
SSE and StarSSE. The first network in SSE or StarSSE is always trained in the same manner as an individual fine-tuned network described in the previous paragraph with the optimal number of fine-tuning epochs and maximum learning rate. We denote these optimal hyperparameter values for each task as  $\times 1$ . For training all the following networks, we use the same learning schedule, for which we vary the number of epochs and maximum learning rate in a cycle. We consider different hyperparameter intervals for different tasks to cover all types of method behavior and denote the values relatively to the optimal ones for individual model fine-tuning (for example,  $\times 0.5$  or  $\times 2$ ).
Aggregating results. All ensemble accuracy values presented in the plots are averaged over several runs, so mean and standard deviation are shown. For Local and Global DEs, we average over at least 5 runs (we fine-tune 5 networks from each pre-trained checkpoint and average over different pre-trainings and fine-tuning random seeds). For SSE and StarSSE, we average over two runs, each from a different pre-trained checkpoint.
Connectivity analysis. To analyze the behavior of ensemble methods, we look at the linear connectivity of ensemble members. Usually, two models are called linearly connected and laying in the same low-loss basin if there is no high-loss barrier on the linear interpolation between them, meaning the loss occurred when linearly connecting the models is lower than the linear interpolation of their losses with the same interpolation coefficients (Entezari et al., 2022). In our analysis, we use accuracy instead of loss. The two metrics behave similarly, however, the absolute accuracy values are easier to interpret. Hence, we examine whether ensemble members lay in the same high-accuracy area (high-accuracy basin) or if there is a zone of lower accuracy on the linear interpolation between them (low-accuracy barrier).
Table 1: Accuracy of ensembles of size 5 of ResNet-50 models with supervised (SV) and self-supervised (SSL) pre-training. Green numbers indicate the gap between Local and Global DE.
Pre-train.	Dataset	Local DE	Global DE	SSE (opt.)	StarSSE (opt.)
SV	CIFAR-100	85.65±0.12	86.82±0.14 (+1.17)	86.04±0.18	86.30±0.18
CIFAR-10	97.52±0.04	97.69±0.05 (+0.17)	97.30±0.05	97.39±0.01
SUN-397	62.23±0.38	64.65±0.10 (+2.42)	63.63±0.14	64.20±0.21
SSL	CIFAR-100	87.33±0.21	88.10±0.11 (+0.77)	87.11±0.06	87.63±0.12
CIFAR-10	97.97±0.05	98.14±0.05 (+0.17)	97.94±0.04	98.06±0.02
SUN-397	65.77±0.19	67.04±0.09 (+1.27)	65.79±0.10	66.37±0.25
Figure 1: Results of SSEs (top row) and StarSEEs (bottom row) of different sizes on CIFAR-100 for different types of pre-training and varying values of cycle hyperparameters (maximum learning rate and number of epochs). Local and Global DEs are shown for comparison with gray dotted lines.
5 Empirical analysis of ensembling in transfer learning
5.1 Local and global deep ensembles
As a first step of our empirical analysis of ensembling methods in transfer learning, we compare the quality of Local and Global DEs in our setup. In Table 1, we show the results for ensembles of maximal considered size 5. Also, the results for ensembles of variable sizes are visualized in Figure 1 with gray dashed lines (see Appendix E for the results on other datasets). For both types of pre-training and all considered target tasks, there is a significant quality gap between Local and Global DEs. Thus, the problem of closing this gap is relevant to our setup.
We also confirm that networks from the Local DE lay in the same high-accuracy basin. Gray baseline lines in Figure 2 show accuracy along line segments between random pairs of networks from the same Local DE (see Appendix F for results on other datasets). There are no significant low-accuracy barriers in training and test metrics.
5.2 Local and semi-local Snapshot Ensembles
Since a Local DE only explores the pre-train basin, we experiment with the SSE method with varying values of cycle hyperparameters to answer two questions: 1) can we improve a Local DE quality by better exploring the pre-train basin, and 2) can we improve it even further by exploring the close vicinity outside the pre-train basin? In this section, we demonstrate the results only on the CIFAR-100 task; the results for other target tasks are similar and presented in Appendix E.
Figure 2: Linear connectivity analysis of Local DEs, SSEs (left plots), and StarSSEs (right plots) on CIFAR-100 with self-supervised pre-training. We show train and test accuracy along line segments between two random networks in Local DEs, between the first and the last (5-th) network in three differently behaving SSEs, and between the first and any other consequent network in three differently behaving StarSSEs. Hyperparameters for more local and more semi-local experiments are the same for SSE and StarSSE, while hyperparameters for the optimal experiments may differ.
Figure 3: Train and test accuracy of individual models from three differently behaving SSEs (left plots) and StarSSEs (right plots) on CIFAR-100 with self-supervised pre-training (with a single fine-tuned model for comparison). Hyperparameters for more local and more semi-local experiments are the same for SSE and StarSSE, while hyperparameters for the optimal experiments may differ.
We vary the maximum learning rate and number of epochs in SSE cycles in the top row of plots in Figure 1. Varying these hyperparameters similarly changes method behavior because training with a higher learning rate or for a larger number of epochs allows the network to go further from the pretrain checkpoint. For low values of hyperparameters, especially the learning rate, SSE collects very closely located checkpoints; therefore, their ensembling does not substantially improve the prediction quality. When we increase the values, SSE collects more diverse checkpoints, and ensemble quality rises. However, for high hyperparameter values, the ensemble quality significantly degrades with the number of networks due to degradation of the quality of individual models as they move further and further away from the pre-train checkpoint (see next sub-sections).
Overall, SSE with optimally chosen hyperparameters can construct an ensemble of similar, and in some cases even better, quality than Local DE but is not able to close the gap with Global DE (see Table 1). Interestingly, SSE shows strong results for small ensembles, sometimes even matches the Global DE for ensembles of size 2, however, for large ensembles, SSE results become weaker (see Figure 1). SSE shows stronger improvements in the case of supervised pre-training, most likely because of the higher difference in the optimal hyperparameter values for SSE and usual Local DEs.
5.3 Analysis of Snapshot Ensemble behavior
To better understand the properties of different types of SSE behavior and the reasons behind them, we analyze three experiments for each setup in more detail. We consider the following three experiments: the more local experiment (with low learning rate and very slow-growing ensemble quality), the more semi-local experiment (with high learning rate, a large number of epochs, and ensemble quality degradation for large ensemble sizes), and the optimal experiment (with medium learning rate and
number of epochs, and the best quality for an ensemble of size 5). See Appendix B for the exact values of hyperparameters for these three experiments. In this section, we demonstrate the results only for experiments with self-supervised pre-training and the CIFAR-100 dataset. The results with supervised pre-training are similar and presented in Appendix F.